from pyquant import get_stock_data, get_stocks_basics, get_db
import pandas as pd
import datetime
import os

# 多因子选股的测试脚本

db = get_db()

# 返回与day最接近的交易日（并且那天可以获取到股票基本面数据）
def get_nearest_open_day(day):
    date = pd.Timestamp(day)
    now = datetime.datetime.now()
    while date < now:
        t = db.stocks.find_one({'code':'sh', 'date':date.strftime("%Y-%m-%d")})
        if t != None:
            t = db.stocks_basics.find_one({'date':date.strftime("%Y-%m-%d")})
            if t != None:
                return date.strftime("%Y-%m-%d")
        date = date + datetime.timedelta(days=1)
    return ""

# 获取到股票从start到end的开盘日期
def get_open_days_count(code, start, end):
    return db.stocks.find({'code': code, 'date': {'$gte': start, '$lte': end}}).count()

# 获取所有的stock列表
def get_stock_list():
    t = list(db.info.find({'name':'stock_list'}))
    t_list = t[0]['value']
    return t_list

# 获取从start到end正常交易的所有股票
def get_open_stocks(start, end):
    stock_list = []
    t_list = get_stock_list()
    days_count = get_open_days_count('sh', start, end)
    for i in t_list:
        t = get_open_days_count(i, start, end)
        if t == days_count:
            stock_list.append(i)

    return stock_list

# 获取date的next month
def next_month(date):
    year = date.year
    month = date.month
    month = month + 1
    if month > 12:
        year = year + 1
        month = 1
    return datetime.date(year, month, 1)

# 提前处理好交易日期
def get_trading_days_pair(start_date, end_date):
    trading_days = []
    t_date = start_date
    while t_date <= end_date:
        t = get_nearest_open_day(t_date.strftime('%Y-%m-%d'))
        if t != "":
            trading_days.append(t)
        t_date = next_month(t_date)

    trading_days_pair = []
    for i in range(len(trading_days) - 1):
        trading_days_pair.append((trading_days[i], trading_days[i + 1]))
    return trading_days_pair

# 获取stock_list列表内的所有股票的那天的基本面数据
def get_selected_stock_basics(date, stock_list):
    data = db.stocks_basics.find({'code':{"$in":stock_list}, "date":date})
    data = pd.DataFrame(list(data))
    if not data.empty:
        del data['_id']
        return data

    return None

# 获取股票那天的收盘价
def get_close_value(code, date):
    data = db.stocks.find_one({'code':code, 'date':date})
    return data['close']

# 在一段交易期，对因子的考察（只需要获取开始那天和结束那天的股票数据， 不需要一天天进行胡测处理）
# 按照“量化投资策略与技术”这本书的多因子章节进行编写
def trading_by_factor(factor, trading_pair, stock_list):
    start = trading_pair[0]
    end = trading_pair[1]
    print(start)
    stock_basics = get_selected_stock_basics(start, stock_list)
    stock_basics = stock_basics.sort_values(by=factor, ascending=True)

    # 对所有股票根据factor排序, 再分组, 每组为其组号
    stock_num = len(stock_list)
    # print(stock_num)
    group_num = 10
    # 未必能获取到所有交易股票的基本面信息!!!
    each_group_num = stock_num // group_num
    groups = {}
    for i in range(group_num):
        groups[i] = []
        t_data = stock_basics[i*each_group_num:(i+1)*each_group_num]
        groups[i].extend(t_data['code'].tolist())

    start_value = get_close_value('sh', start)
    end_value = get_close_value('sh', end)
    sh_change = (end_value - start_value) / start_value
    groups_change = {}
    # 跑赢市场
    groups_win_count = {}
    groups_win_rate = {}
    groups_change_avg = {}
    for i in range(group_num):
        groups_change[i] = []
        groups_win_count[i] = 0
        t_sum = 0
        for each in groups[i]:
            start_value = get_close_value(each, start)
            end_value = get_close_value(each, end)
            change = (end_value - start_value) / start_value
            groups_change[i].append(change)
            if change > sh_change:
                groups_win_count[i] = groups_win_count[i] + 1
            t_sum += change
        groups_win_rate[i] = groups_win_count[i] / each_group_num
        groups_change_avg[i] = t_sum / each_group_num

    change_avg = 0
    id_avg = 0
    for i in range(group_num):
        change_avg += groups_change_avg[i]
        id_avg += i
    change_avg = change_avg / group_num
    id_avg = id_avg / group_num

    covariance = 0
    for i in range(group_num):
        covariance += (groups_change_avg[i] - change_avg) * (i - id_avg)
    covariance /= group_num
    # check covariance 需要大于一个阈值!!
    print("covariance = " + str(covariance))

    # check 中间的要介于两边
    print("======")
    for i in range(group_num):
        print(groups_change_avg[i])
    print("======")
    for i in range(group_num):
        print(groups_win_rate[i])
    print("======")

    # check 一个要远高于市场, 一个远低于市场
    print(str(change_avg) + "   " + str(groups_change_avg[0]) + "  " + str(groups_change_avg[group_num - 1]))

    print('----------------------------------')
    # print(stock_basics)

# 因为基本面数据太少了， 目前最多只有这么多数据
def check_factor(factor, start='2016-09-01', end='2017-08-01'):
    # 开始日期和结束日期， 转化成timestamp格式， 并把它们换成月初
    start_date = pd.Timestamp(start)
    end_date = pd.Timestamp(end)

    t_year = start_date.year
    t_month = start_date.month
    start_date = datetime.date(t_year, t_month, 1)
    start = start_date.strftime('%Y-%m-%d')

    t_year = end_date.year
    t_month = end_date.month
    end_date = datetime.date(t_year, t_month, 1)
    end = end_date.strftime('%Y-%m-%d')
    # print(start + "  " + end)

    # 获取从start到end这段时间内正常交易的股票
    stock_list = get_open_stocks(start, end)

    # 提前处理好交易时间的pair
    trading_days_pair = get_trading_days_pair(start_date, end_date)
    for each in trading_days_pair:
        trading_by_factor(factor, each, stock_list)


# 测试pe是否是好的因子
check_factor('pe')